{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Figure 5: sphere tracing steps visualization\"\"\"\n",
    "\n",
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "import os\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "from disk_tracing import *\n",
    "\n",
    "import mitsuba as mi\n",
    "\n",
    "from sdf2d.shapes import Grid2d, UnionSDF\n",
    "\n",
    "import skfmm\n",
    "\n",
    "sdf_res = 512\n",
    "y, x = dr.meshgrid(*[dr.linspace(mi.Float, 0, 1, sdf_res) for i in range(2)], indexing='ij')\n",
    "x = x * 0.9\n",
    "mask = y > (dr.sin(10*x - 4) * x*0.7 + 0.9 - 0.7 * dr.sin(x))\n",
    "mask = mask | (x > 0.8) \n",
    "mask = np.flipud(np.reshape(mask, (sdf_res, sdf_res)))\n",
    "sdf = Grid2d(skfmm.distance(0.5 - mask, 1 / sdf_res))\n",
    "\n",
    "fig_name = 'diff_sphere_tracing'\n",
    "\n",
    "def eval_warp_field(sdf, res, eval_pos=None):\n",
    "    if eval_pos is None:\n",
    "        y_v, x_v = dr.meshgrid(*[dr.linspace(mi.Float, 0 + 0.5/res, 1-0.5/res, res) for i in range(2)], indexing='ij')\n",
    "        p_v = mi.Vector2f(x_v, y_v)\n",
    "    else:\n",
    "        p_v = eval_pos \n",
    "        x_v = eval_pos.x\n",
    "        y_v = eval_pos.y\n",
    "\n",
    "    sdf_grad = sdf.eval_grad(p_v)\n",
    "    sdf_value = sdf.eval(p_v)\n",
    "    warp = -sdf_value * dr.detach(sdf_grad / dr.squared_norm(sdf_grad))\n",
    "    return x_v, y_v, warp\n",
    "\n",
    "RAY_COLOR = [0.2, 0.3, 0.5]\n",
    "WARP_T_COLOR = [0.8, 0.2, 0.2]\n",
    "SURFACE_COLOR = [1.0, 0.9, 0.7]\n",
    "BG_COLOR = [0.95, 0.95, 0.95]\n",
    "draw_sdf = False\n",
    "\n",
    "y_min = 0.02\n",
    "y_max = 0.85\n",
    "\n",
    "y_min = 0.0\n",
    "y_max = 0.9\n",
    "\n",
    "n_isolines = 33\n",
    "y_offset = -0.13\n",
    "fontsize = 12\n",
    "base_size = 4\n",
    "n_rows = 1\n",
    "n_cols = 4\n",
    "\n",
    "total_width = TEXT_WIDTH\n",
    "aspect = n_rows / n_cols * 1.02\n",
    "\n",
    "fig = plt.figure(1, figsize=(total_width, total_width * aspect), constrained_layout=False)\n",
    "gs = fig.add_gridspec(n_rows, n_cols, wspace=0.025, hspace=0.025)\n",
    "\n",
    "# ---------------------------- Plot 0 ------------------------------------\n",
    "# Plot SDF and its isolines\n",
    "ax = fig.add_subplot(gs[0])\n",
    "\n",
    "res = 1024\n",
    "vector_field_res = 12\n",
    "y, x = dr.meshgrid(*[dr.linspace(mi.Float, 0, 1, res) for i in range(2)], indexing='ij')\n",
    "p = mi.Vector2f(x, y)\n",
    "\n",
    "sdf_values = sdf.eval(p)\n",
    "r = 0.6\n",
    "ax.imshow(np.reshape(sdf_values, (res, res)), interpolation='none', extent=[0,1,0,1], cmap='coolwarm', vmin=-r, vmax=r, origin='lower')\n",
    "ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), levels=[0], colors='k')\n",
    "ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), \n",
    "                    levels=np.linspace(-1,1,n_isolines), alpha=0.9, colors='k', linewidths=0.2)\n",
    "disable_ticks(ax)\n",
    "x_v, y_v, warp = eval_warp_field(sdf, vector_field_res)\n",
    "ax.quiver(x_v, y_v, warp.x, warp.y, scale=4)\n",
    "ax.set_title(r\"(a) $\\mathcal{V}(\\mathbf{x}, \\bm{\\pi})$\", y=y_offset, fontsize=fontsize)\n",
    "ax.set_ylim(y_min, y_max)\n",
    "ax.set_xlim(0, 1)\n",
    "\n",
    "# ---------------------------- Plot 1 ------------------------------------\n",
    "# Plot SDF vector field\n",
    "ax = fig.add_subplot(gs[1])\n",
    "\n",
    "if type(sdf) is UnionSDF:\n",
    "    dr.enable_grad(sdf.sdf1.p.y)\n",
    "    dr.set_grad(sdf.sdf1.p.y, 0.0)\n",
    "else:\n",
    "    dr.enable_grad(sdf.p.y)\n",
    "    dr.set_grad(sdf.p.y, 0.0)\n",
    "\n",
    "x_v, y_v, warp = eval_warp_field(sdf, vector_field_res)\n",
    "\n",
    "if type(sdf) is UnionSDF:\n",
    "    dr.forward(sdf.sdf1.p.y)\n",
    "else:\n",
    "    dr.forward(sdf.p.y)\n",
    "\n",
    "warp = dr.grad(warp)\n",
    "\n",
    "y, x = dr.meshgrid(*[dr.linspace(mi.Float, 0, 1, res) for i in range(2)], indexing='ij')\n",
    "p = mi.Vector2f(x, y)\n",
    "\n",
    "sdf_values = sdf.eval(p)\n",
    "r = 0.6\n",
    "ax.imshow(np.reshape(sdf_values, (res, res)), interpolation='none', extent=[0,1,0,1], cmap='coolwarm', vmin=-r, vmax=r, origin='lower')\n",
    "ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), levels=[0], colors='k')\n",
    "ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), \n",
    "                    levels=np.linspace(-1,1,n_isolines), alpha=0.9, colors='k', linewidths=0.2)\n",
    "\n",
    "mask = np.zeros((vector_field_res, vector_field_res))\n",
    "mask[2:-6, -1:] = 1\n",
    "mask = 1 - mask\n",
    "mask = mask.astype(bool).ravel()\n",
    "\n",
    "x_v = np.array(x_v)[mask]\n",
    "y_v = np.array(y_v)[mask]\n",
    "warp_x = np.array(warp.x)[mask]\n",
    "warp_y = np.array(warp.y)[mask]\n",
    "\n",
    "ax.quiver(x_v, y_v, warp_x, warp_y)\n",
    "\n",
    "ax.arrow(0.93, 0.2, 0.0, 0.3, head_width=0.03, overhang=0.2, facecolor='k', length_includes_head=True, zorder=10)\n",
    "txt = ax.text(0.94, 0.35, r'$\\pi$')\n",
    "txt.set_path_effects([path_effects.withStroke(linewidth=1.0, foreground='white')])\n",
    "\n",
    "ax.set_title(r\"(b) $\\partial_{\\pi}\\mathcal{V}(\\mathbf{x}, \\bm{\\pi})$\", y=y_offset, fontsize=fontsize)\n",
    "disable_ticks(ax)\n",
    "ax.set_ylim(y_min, y_max)\n",
    "ax.set_xlim(0, 1)\n",
    "# ---------------------------- Plot 2 ------------------------------------\n",
    "ax = fig.add_subplot(gs[2])\n",
    "ray = mi.Ray2f([0.0, 0.4], [1, 0])\n",
    "ray = mi.Ray2f([0.0, 0.6], [1, 0])\n",
    "its_t, warp_t_integral, points, dists, _ = intersect_sdf_simple(sdf, ray)\n",
    "points = np.array(points)\n",
    "dists = np.array(dists)\n",
    "r = 35000\n",
    "use_log_space = False\n",
    "sdf_values = sdf.eval(p)\n",
    "sdf_grad_values = sdf.eval_grad(p)\n",
    "\n",
    "weights = sphere_tracing_step_weight(ray, sdf_values, sdf_grad_values)\n",
    "if use_log_space:\n",
    "    weights = np.log(weights)\n",
    "    r = 8\n",
    "\n",
    "im_weights = ax.imshow(np.reshape(weights, (res, res)), interpolation='none', extent=[0,1,0,1], cmap='coolwarm', vmin=0, vmax=r, origin='lower')\n",
    "ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), levels=[0], colors='k')\n",
    "ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), \n",
    "                    levels=np.linspace(-1,1,3), alpha=0.9, colors='k', linewidths=0.2, zorder=10)\n",
    "\n",
    "plt.scatter(points[:, 0], points[:, 1], color='red',zorder=30)\n",
    "for p, r in zip(points, dists):\n",
    "    ptc = matplotlib.patches.Circle(p, r, facecolor=[1,1,1], alpha=0.1, lw=0.5, edgecolor=None, zorder=20)\n",
    "    ax.add_patch(ptc)\n",
    "    ptc = matplotlib.patches.Circle(p, r, fill=None, alpha=1.0, lw=0.5, edgecolor='white', zorder=22)\n",
    "    ax.add_patch(ptc)\n",
    "\n",
    "ray_o = np.array(ray.o).ravel()\n",
    "ray_d = np.array(ray.d).ravel() * np.array(its_t).ravel()\n",
    "ax.arrow(ray_o[0], ray_o[1], ray_d[0], ray_d[1], head_width=0.03, overhang=0.2, facecolor='k', length_includes_head=True)\n",
    "ax.set_ylim(y_min, y_max)\n",
    "ax.set_xlim(0, 1)\n",
    "disable_ticks(ax)\n",
    "ax.set_title(\"(c) Sphere tracing steps\", y=y_offset, fontsize=fontsize)\n",
    "\n",
    "cbar_ax = fig.add_axes([0.652, 0.23, 0.05, 0.01])\n",
    "cb = fig.colorbar(im_weights, cax=cbar_ax, label='Weight', orientation='horizontal')\n",
    "cb.set_label(label='Weight', size=fontsize)\n",
    "cbar_ax.xaxis.set_ticks_position('top')\n",
    "xlbl = cbar_ax.xaxis.get_label()\n",
    "xlbl.set_color('white')\n",
    "cbar_ax.tick_params(axis='x', colors='white', labelsize=9)\n",
    "\n",
    "# ---------------------------- Plot 3 ------------------------------------\n",
    "ax = fig.add_subplot(gs[3])\n",
    "n_rays = 128\n",
    "ray_d = mi.Vector2f(dr.zeros(mi.Float, n_rays) + 1, dr.zeros(mi.Float, n_rays))\n",
    "ray = mi.Ray2f(mi.Point2f(0.0, dr.linspace(mi.Float, 0, 1, n_rays)), ray_d)\n",
    "its_t, warp_t_integral, points, dists, _ = intersect_sdf_simple(sdf, ray)\n",
    "\n",
    "r = 0.6\n",
    "if draw_sdf:\n",
    "    ax.imshow(np.reshape(sdf_values, (res, res)), interpolation='none', extent=[0,1,0,1], cmap='coolwarm', vmin=-r, vmax=r, origin='lower')\n",
    "    ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), \n",
    "                    levels=np.linspace(-1,1,n_isolines), alpha=0.9, colors='k', linewidths=0.2, zorder=10)\n",
    "else:\n",
    "    colors = dr.select(sdf_values < 0, mi.Vector3f(SURFACE_COLOR), mi.Vector3f(BG_COLOR))\n",
    "    ax.imshow(np.reshape(colors, (res, res, 3)), interpolation='none', extent=[0,1,0,1], origin='lower')\n",
    "\n",
    "ray_o_values = np.array(ray.o.y)\n",
    "ax.plot(warp_t_integral, ray_o_values, color=[0.8, 0.2, 0.2], lw=3)\n",
    "\n",
    "# Compute actual warp field on the eval positions\n",
    "if False:\n",
    "    dr.enable_grad(sdf.p.y)\n",
    "    dr.set_grad(sdf.p.y, 0.0)\n",
    "    x_v, y_v, warp = eval_warp_field(sdf, vector_field_res, eval_pos=dr.detach(ray(warp_t_integral)))\n",
    "    dr.forward(sdf.p.y)\n",
    "    warp = dr.grad(warp)\n",
    "    freq = 19\n",
    "    x_v = np.array(x_v)[::freq]\n",
    "    y_v = np.array(y_v)[::freq]\n",
    "    warp_x = np.array(warp.x)[::freq]\n",
    "    warp_y = np.array(warp.y)[::freq]\n",
    "    ax.quiver(x_v, y_v, warp_x, warp_y, zorder=20)\n",
    "\n",
    "ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), levels=[0], colors='k')\n",
    "\n",
    "# Evaluate some fewer rays to actually draw\n",
    "n_rays = 8\n",
    "ray_d = mi.Vector2f(dr.zeros(mi.Float, n_rays) + 1, dr.zeros(mi.Float, n_rays))\n",
    "ray = mi.Ray2f(mi.Point2f(0.0, dr.linspace(mi.Float, 0 + 0.5 / n_rays, 1 - 0.5 / n_rays, n_rays)), ray_d)\n",
    "its_t, warp_t_integral, points, dists, _ = intersect_sdf_simple(sdf, ray)\n",
    "\n",
    "for ray_o, ray_d, dist in zip(np.array(ray.o), np.array(ray.d), np.array(warp_t_integral)):\n",
    "    ray_d = ray_d * dist\n",
    "    ax.arrow(ray_o[0], ray_o[1], ray_d[0], ray_d[1], head_width=0.03, overhang=0.2, edgecolor=RAY_COLOR, facecolor=RAY_COLOR, length_includes_head=True, zorder=10)\n",
    "ax.set_ylim(y_min, y_max)\n",
    "ax.set_xlim(0, 1)\n",
    "disable_ticks(ax)\n",
    "ax.set_title(\"(d) Evaluation distance\", y=y_offset, fontsize=fontsize)\n",
    "\n",
    "plt.margins(0, 0)\n",
    "# save_fig(fig_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
